﻿// --- ProcessExtensions.cs (Updated) ---
using System;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Security.Principal;
using System.Threading;
using System.Threading.Tasks;

namespace DefenderUI
{
    public static class ProcessExtensions
    {
        // --- Win32 API Imports for getting process owner ---
        [DllImport("advapi32.dll", SetLastError = true)]
        private static extern bool OpenProcessToken(IntPtr ProcessHandle, uint DesiredAccess, out IntPtr TokenHandle);

        [DllImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        private static extern bool CloseHandle(IntPtr hObject);

        /// <summary>
        /// Gets the owner (user account) of a specified process.
        /// </summary>
        /// <param name="process">The process for which to find the owner.</param>
        /// <returns>The user account name or an error message if the owner cannot be retrieved.</returns>
        public static string GetProcessOwner(this Process process)
        {
            IntPtr processHandle = IntPtr.Zero;
            try
            {
                // Opens the security token associated with the process.
                if (!OpenProcessToken(process.Handle, 0x0008, out processHandle))
                {
                    // Access might be denied for system-level processes.
                    return "N/A";
                }

                // Use the token to get the WindowsIdentity, which represents the user.
                using (var wi = new WindowsIdentity(processHandle))
                {
                    // Return the user's name (e.g., "SYSTEM", "MyPC\Username").
                    return wi.Name;
                }
            }
            catch
            {
                // Catch any other exceptions (e.g., process has just exited).
                return "N/A";
            }
            finally
            {
                // Crucially, close the handle to the process token to prevent resource leaks.
                if (processHandle != IntPtr.Zero)
                {
                    CloseHandle(processHandle);
                }
            }
        }

        /// <summary>
        /// Asynchronously waits for a process to exit.
        /// </summary>
        public static Task WaitForExitAsync(this Process process, CancellationToken cancellationToken = default)
        {
            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
            process.EnableRaisingEvents = true;
            process.Exited += (sender, args) => tcs.TrySetResult(null);

            if (cancellationToken != default)
            {
                cancellationToken.Register(() => {
                    try { process.Kill(); } catch { /* Ignore */ }
                    tcs.TrySetCanceled();
                });
            }
            return tcs.Task;
        }
    }
}